# -*- coding: utf-8 -*-
"""
@Time    : 2025/3/11 11:14 
@Author  : ZhangShenao 
@File    : 2.Bert分词器详解.py
@Desc    : 使用BERT的分词器进行编解码
"""

from transformers import BertTokenizer

# 加载Bert的分词器
cache_dir = "./models/bert-base-chinese/models--bert-base-chinese/snapshots/c30a6ed22ab4564dc1e3b2ecbf6e766b0611a33f"
tokenizer = BertTokenizer.from_pretrained(cache_dir)
print(f"Bert分词器加载完成: \n{tokenizer}")

# 准备要编码的文本数据
sentences = ["明月几时有，",
             "君不见黄河之水天上来，"]

# 使用tokenizer批量编码
# 将文本转化为向量(仅仅是编码处理,而不是Embedding过程)
# 在开启Padding处理之后,编码结果会按照max_length长度对齐: 超过的部分会截断,不足的部分会填充0
encoded = tokenizer.batch_encode_plus(
    batch_text_or_text_pairs=sentences,  # 指定原始文本
    add_special_tokens=True,  # 添加特殊的Token(如起始符、终止符等)
    # 当句子长度大于max_length(上限是model_max_length)时，截断
    max_length=10,  # 指定编码序列的最大长度,上限为tokenizer_config文件中配置的model_max_length
    truncation=True,  # 开启截断:当编码后的序列长度大于max_length时,则触发截断
    padding="max_length",  # 按照max_length最大长度进行填充,当序列长度不足时,一律填充0
    return_tensors=None,  # 返回的编码序列格式:可取值为tf,pt,np,默认为list
    return_attention_mask=True,  # 返回attention mask注意力掩码
    return_token_type_ids=True,  # 返回token type ids
    return_special_tokens_mask=True,  # 返回特殊Token的掩码
    return_length=True  # 返回序列长度
)

print(f"编码结果: \n{encoded}")

# 按照key-value逐个打印编码结果
# input_ids: 编码结果,即编码后的Token ID向量
# token_type_ids: 仅针对上下文编码的场景(一问一答): 第一个句子和特殊符号的位置值为0,第二个句子的位置值为1
# special_tokens_mask: 特殊符号的位置是1,其他位置是0
# length: 编码之后的序列长度
print("编码结果详情:")
for k, v in encoded.items():
    print(f"{k}: {v}")

# 使用tokenizer进行解码
# 查看具体内容
print(f"第一句解码结果: \n{tokenizer.decode(encoded['input_ids'][0])}")
print(f"第二句解码结果: \n{tokenizer.decode(encoded['input_ids'][1])}")
